import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont


DISTINCT_COLORS = ['#e6194b', '#3cb44b', '#ffe119', '#0082c8', '#f58231', '#911eb4', '#46f0f0', 
                   '#f032e6', '#d2f53c', '#fabebe', '#008080', '#000080', '#aa6e28', '#fffac8', 
                   '#800000', '#aaffc3', '#808000', '#ffd8b1', '#e6beff', '#808080', '#FFFFFF']


def get_color_map(labels):
    """
    use repeated colors if len(labels) > 21
    """
    label_color_map = {}
    for idx, label in enumerate(labels):
        color = DISTINCT_COLORS[idx % len(DISTINCT_COLORS)]
        label_color_map[label] = color
    return label_color_map


def view_data(image, boxes, labels):
    """
    helper function to visualize the effect of transformation utilities.
    
    """
    width = image.size[0]
    height = image.size[1]

    cv_image = np.array(image)
    cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
    
    np_boxes = boxes.clone().detach().cpu().numpy()
    nP_labels = labels.clone().detach().cpu().numpy()

    np_boxes[:, [0, 2]] *= width
    np_boxes[:, [1, 3]] *= height
    np_boxes = np_boxes.astype(np.int32)
    
    for i, box in enumerate(np_boxes):
        xmin = box[0]
        ymin = box[1]
        xmax = box[2]
        ymax = box[3]
        cv2.rectangle(cv_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
    
    cv2.imshow('image', cv_image)
    cv2.waitKey(0)


def get_boxes_by_drawing(width, height):
    """
    draw boxes on a window using opencv
    """
    raw_boxes = []
    boxes = []
    canvas = np.zeros((height, width, 3), dtype=np.uint8)
    while True:
        # draw raw boxes.
        canvas1 = np.copy(canvas)
        for b in raw_boxes:
            x, y, w, h = b
            cv2.rectangle(canvas1, (x, y), (x+w, y+h), (0, 255, 0), 1)
        
        box = cv2.selectROI('draw box', canvas1, False, False)
        print(box)
        if box[2] < 10 and box[3] < 10:
            break
        raw_boxes.append(box)

        xmin, ymin, w, h = box
        boxes.append([xmin / width, ymin / height, (xmin + w) / width, (ymin + h) / height])
    return boxes


def visualize_detections_pil(image, boxes, labels, color_map=None):
    """
    draw detected boxes on original image
    boxes are torch.Tensor of absolute coordinates
    labels are a list of class names

    # TODO: support boxes of both type np.ndarray or torch.Tensor.
    """

    if color_map is None:
        # use default color '#00FF00' (green)
        color = '#00FF00'

    annotated_image = image
    draw = ImageDraw.Draw(annotated_image)
    font = ImageFont.truetype('./calibril.ttf', 15)

    for i in range(boxes.shape[0]):
        box_location = boxes[i]
        label = labels[i]

        if color_map is None:
            color = '#00FF00'
        else:
            color = color_map[label]
        
        # draw box
        draw.rectangle(xy=[l for l in box_location], outline=color)
        draw.rectangle(xy=[l+1. for l in box_location], outline=color)

        # draw text
        text_size = font.getsize(label.upper())
        text_location = [box_location[0] + 1, box_location[1] - text_size[1]]
        textbox_location = [box_location[0], box_location[1] - text_size[1], 
                            box_location[0] + text_size[0] + 4, box_location[1]]
        draw.rectangle(xy=textbox_location, fill=color)
        draw.text(xy=text_location, text=label.upper(), fill='white', font=font)
    del draw
    return annotated_image
